{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "700aec2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/pt/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#测试数据\n",
    "targets = [{\n",
    "    'class_labels': torch.randint(low=0, high=4, size=[5]),\n",
    "    'boxes': torch.rand(5, 4),\n",
    "} for _ in range(8)]\n",
    "\n",
    "outputs = {\n",
    "    'logits': torch.randn(8, 100, 92),\n",
    "    'pred_boxes': torch.rand(8, 100, 4)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "309a460f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([ 7,  8, 35, 92, 97]), tensor([0, 1, 4, 2, 3])),\n",
       " (tensor([35, 36, 44, 59, 61]), tensor([2, 1, 0, 4, 3])),\n",
       " (tensor([10, 31, 37, 54, 63]), tensor([4, 0, 2, 3, 1])),\n",
       " (tensor([ 2, 21, 39, 55, 84]), tensor([3, 2, 1, 4, 0])),\n",
       " (tensor([12, 15, 19, 20, 23]), tensor([2, 0, 3, 1, 4])),\n",
       " (tensor([ 7, 48, 58, 63, 87]), tensor([0, 4, 2, 1, 3])),\n",
       " (tensor([13, 40, 60, 85, 90]), tensor([0, 4, 1, 2, 3])),\n",
       " (tensor([15, 23, 35, 68, 83]), tensor([3, 1, 2, 0, 4]))]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy.optimize._lsap import linear_sum_assignment\n",
    "\n",
    "\n",
    "class DetrHungarianMatcher:\n",
    "\n",
    "    def box_iou(self, boxes1, boxes2):\n",
    "        inter = []\n",
    "        union = []\n",
    "        area = []\n",
    "        for box1 in boxes1:\n",
    "            for box2 in boxes2:\n",
    "                #求交集面积\n",
    "                x1 = torch.max(box1[0], box2[0])\n",
    "                y1 = torch.max(box1[1], box2[1])\n",
    "                x2 = torch.min(box1[2], box2[2])\n",
    "                y2 = torch.min(box1[3], box2[3])\n",
    "\n",
    "                w = (x2 - x1).clamp(min=0)\n",
    "                h = (y2 - y1).clamp(min=0)\n",
    "\n",
    "                inter.append(w * h)\n",
    "\n",
    "                #求并集面积\n",
    "                w1 = box1[2] - box1[0]\n",
    "                h1 = box1[3] - box1[1]\n",
    "                w2 = box2[2] - box2[0]\n",
    "                h2 = box2[3] - box2[1]\n",
    "\n",
    "                s1 = w1 * h1\n",
    "                s2 = w2 * h2\n",
    "\n",
    "                union.append(s1 + s2 - (w * h))\n",
    "\n",
    "                #求扩展面积\n",
    "                x1 = torch.min(box1[0], box2[0])\n",
    "                y1 = torch.min(box1[1], box2[1])\n",
    "                x2 = torch.max(box1[2], box2[2])\n",
    "                y2 = torch.max(box1[3], box2[3])\n",
    "\n",
    "                w = (x2 - x1).clamp(min=0)\n",
    "                h = (y2 - y1).clamp(min=0)\n",
    "\n",
    "                area.append(w * h)\n",
    "\n",
    "        inter = torch.stack(inter).reshape(len(boxes1), len(boxes2))\n",
    "        union = torch.stack(union).reshape(len(boxes1), len(boxes2))\n",
    "        area = torch.stack(area).reshape(len(boxes1), len(boxes2))\n",
    "\n",
    "        #前面的数是iou,值域0-1,衡量了两个框重合的程度,这个数越大越好\n",
    "\n",
    "        #后面的数是个分数,分别来看分子和分母\n",
    "        #分子是扩展面积-并集面积,这个数显然是越小越好.\n",
    "        #分母是扩展面积,显然是起归一化作用,所以这个分数的值域是0-1\n",
    "\n",
    "        #综合以上,这个数总的来说还是iou,只是额外考虑了扩展面积的情况\n",
    "        return (inter / union) - (area - union) / area\n",
    "\n",
    "    #等价写法,上面的写法效率低\n",
    "    def box_iou(self, boxes1, boxes2):\n",
    "        area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])\n",
    "        area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])\n",
    "\n",
    "        p1 = torch.max(boxes1[:, :2].unsqueeze(1), boxes2[:, :2])\n",
    "        p2 = torch.min(boxes1[:, 2:].unsqueeze(1), boxes2[:, 2:])\n",
    "        wh = (p2 - p1).clamp(min=0)\n",
    "        inter = wh[:, :, 0] * wh[:, :, 1]\n",
    "\n",
    "        union = area1.unsqueeze(1) + area2 - inter\n",
    "\n",
    "        p1 = torch.min(boxes1[:, :2].unsqueeze(1), boxes2[:, :2])\n",
    "        p2 = torch.max(boxes1[:, 2:].unsqueeze(1), boxes2[:, 2:])\n",
    "        wh = (p2 - p1).clamp(min=0)\n",
    "        area = wh[:, :, 0] * wh[:, :, 1]\n",
    "\n",
    "        return (inter / union) - (area - union) / area\n",
    "\n",
    "    def xywh_to_x1y1x2y2(self, boxes):\n",
    "        x = boxes[:, 0]\n",
    "        y = boxes[:, 1]\n",
    "        w = boxes[:, 2]\n",
    "        h = boxes[:, 3]\n",
    "\n",
    "        x1 = x - 0.5 * w\n",
    "        y1 = y - 0.5 * h\n",
    "        x2 = x + 0.5 * w\n",
    "        y2 = y + 0.5 * h\n",
    "\n",
    "        return torch.stack([x1, y1, x2, y2], dim=-1)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def __call__(self, outputs, targets):\n",
    "        #取所有框和预测结果\n",
    "        #[8, 100, 92] -> [800, 92]\n",
    "        logits = outputs['logits'].flatten(0, 1).softmax(1)\n",
    "\n",
    "        #[8, 100, 4] -> [800, 4]\n",
    "        pred_boxes = outputs['pred_boxes'].flatten(0, 1)\n",
    "\n",
    "        #取目标\n",
    "        #[52]\n",
    "        class_labels = torch.cat([i['class_labels'] for i in targets])\n",
    "\n",
    "        #[52, 4]\n",
    "        target_boxes = torch.cat([i['boxes'] for i in targets])\n",
    "\n",
    "        #label的loss,简单的预测概率取反\n",
    "        #[800, 92] -> [800, 52]\n",
    "        class_cost = -logits[:, class_labels]\n",
    "\n",
    "        #框的loss,4个点距离的和作为loss\n",
    "        #[800, 4],[52, 4] -> [800, 52]\n",
    "        bbox_cost = []\n",
    "        for box1 in pred_boxes:\n",
    "            cost = [(box1 - box2).abs().sum() for box2 in target_boxes]\n",
    "            bbox_cost.append(torch.stack(cost))\n",
    "        bbox_cost = torch.stack(bbox_cost)\n",
    "\n",
    "        #等价写法\n",
    "        #bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)\n",
    "\n",
    "        #[800, 52]\n",
    "        giou_cost = -self.box_iou(self.xywh_to_x1y1x2y2(pred_boxes),\n",
    "                                  self.xywh_to_x1y1x2y2(target_boxes))\n",
    "\n",
    "        #[800, 52] -> [8, 100, 52]\n",
    "        cost = 5 * bbox_cost + 1 * class_cost + 2 * giou_cost\n",
    "        cost = cost.view(8, 100, -1).cpu()\n",
    "\n",
    "        indices = []\n",
    "        sum_s = 0\n",
    "        for c, t in zip(cost, targets):\n",
    "            #取目标框的数量\n",
    "            s = len(t['boxes'])\n",
    "\n",
    "            #取这些框的计算结果\n",
    "            #[100, lens]\n",
    "            c = c[:, sum_s:sum_s + s]\n",
    "\n",
    "            #累计索引\n",
    "            sum_s = sum_s + s\n",
    "\n",
    "            #c这个矩阵记录的是loss的情况\n",
    "            #求最小loss的分配方式\n",
    "            index_row, index_col = linear_sum_assignment(c)\n",
    "            index_row = torch.LongTensor(index_row)\n",
    "            index_col = torch.LongTensor(index_col)\n",
    "            indices.append((index_row, index_col))\n",
    "\n",
    "        return indices\n",
    "\n",
    "\n",
    "matcher = DetrHungarianMatcher()\n",
    "\n",
    "matcher(outputs, targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5eceb918",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss_ce': tensor(4.9638),\n",
       " 'loss_bbox': tensor(0.4050),\n",
       " 'loss_giou': tensor(0.7186)}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class DetrLoss(torch.nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        empty_weight = torch.ones(92)\n",
    "        empty_weight[-1] = 0.1\n",
    "        self.register_buffer('empty_weight', empty_weight)\n",
    "\n",
    "    def loss_labels(self, outputs, targets, indices):\n",
    "        # 默认都是背景\n",
    "        #[8, 100]\n",
    "        target_classes = torch.full([8, 100],\n",
    "                                    91,\n",
    "                                    dtype=torch.int64,\n",
    "                                    device=outputs['logits'].device)\n",
    "\n",
    "        #遍历8条数据\n",
    "        for i in range(8):\n",
    "            #遍历每一个分配结果(最小cost方式分配)\n",
    "            for io, it in zip(*indices[i]):\n",
    "                #按照最小cost的方式分配每个预测结果的目标\n",
    "                target_classes[i, io.item()] = targets[i]['class_labels'][\n",
    "                    it.item()]\n",
    "\n",
    "        #[8, 100, 92] -> [8, 92, 100]\n",
    "        logits = outputs['logits'].transpose(1, 2)\n",
    "\n",
    "        return torch.nn.functional.cross_entropy(logits, target_classes,\n",
    "                                                 self.empty_weight)\n",
    "\n",
    "    def loss_boxes(self, outputs, targets, indices):\n",
    "        boxes_output = []\n",
    "        boxes_target = []\n",
    "        # 遍历8条数据\n",
    "        for i in range(8):\n",
    "            # 遍历每一个分配结果(最小cost方式分配)\n",
    "            for io, it in zip(*indices[i]):\n",
    "                # 按照最小cost的方式取每一对框\n",
    "                boxes_output.append(outputs['pred_boxes'][i, io.item()])\n",
    "                boxes_target.append(targets[i]['boxes'][it.item()])\n",
    "\n",
    "        boxes_output = torch.stack(boxes_output)\n",
    "        boxes_target = torch.stack(boxes_target)\n",
    "\n",
    "        num_boxes = sum(len(i['class_labels']) for i in targets)\n",
    "        if num_boxes < 1:\n",
    "            num_boxes = 1\n",
    "\n",
    "        #没对框之间求L1距离作为loss\n",
    "        loss_bbox = torch.nn.functional.l1_loss(boxes_output,\n",
    "                                                boxes_target,\n",
    "                                                reduction='none')\n",
    "        loss_bbox = loss_bbox.sum() / num_boxes\n",
    "\n",
    "        #取iou作为第二部分的loss\n",
    "        #只需要考虑成对的框的iou,所以取对角线元素计算即可\n",
    "        #iou是越大越好,优化方向取反,所以loss=1-iou\n",
    "        giou = matcher.box_iou(matcher.xywh_to_x1y1x2y2(boxes_output),\n",
    "                               matcher.xywh_to_x1y1x2y2(boxes_target))\n",
    "        loss_giou = 1 - giou.diag()\n",
    "        loss_giou = loss_giou.sum() / num_boxes\n",
    "\n",
    "        return loss_bbox, loss_giou\n",
    "\n",
    "    def forward(self, outputs, targets):\n",
    "        indices = matcher(outputs, targets)\n",
    "\n",
    "        losses = {}\n",
    "        losses['loss_ce'] = self.loss_labels(outputs, targets, indices)\n",
    "\n",
    "        loss_bbox, loss_giou = self.loss_boxes(outputs, targets, indices)\n",
    "        losses['loss_bbox'] = loss_bbox\n",
    "        losses['loss_giou'] = loss_giou\n",
    "\n",
    "        return losses\n",
    "\n",
    "\n",
    "criterion = DetrLoss()\n",
    "\n",
    "criterion(outputs, targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7820caf3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss_ce': tensor(4.9638),\n",
       " 'loss_bbox': tensor(0.4050),\n",
       " 'loss_giou': tensor(0.7186),\n",
       " 'cardinality_error': tensor(94.1250)}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def test():\n",
    "    from transformers.models.detr.modeling_detr import DetrLoss, DetrHungarianMatcher\n",
    "\n",
    "    matcher = DetrHungarianMatcher(class_cost=1, bbox_cost=5, giou_cost=2)\n",
    "\n",
    "    criterion = DetrLoss(matcher=matcher,\n",
    "                         num_classes=91,\n",
    "                         eos_coef=0.1,\n",
    "                         losses=['labels', 'boxes', 'cardinality'])\n",
    "\n",
    "    return criterion(outputs, targets)\n",
    "\n",
    "\n",
    "test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt]",
   "language": "python",
   "name": "conda-env-pt-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
